package interceptor

import (
	"context"
	"gozero-microservices/user-service/rpc/internal/svc"
	"gozero-microservices/user-service/rpc/utils"

	"github.com/zeromicro/go-zero/core/limit"
	"github.com/zeromicro/go-zero/core/stores/redis"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

/**
 * 滑动窗口限流器
 * 	1.在固定时间窗口内统计请求次数，超限则拒绝
 * 	2.需要强行限制瞬时请求量，例如防止数据库被刷
 * 	3.固定窗口，规则严格
 * 	4.如果你的场景需要简单粗暴地在单位时间内封顶请求数，比如某个关键接口的严格QPS上限，滑动窗口 (periodlimit) 很合适
 *
 * @param *svc.ServiceContext svcCtx -------------------------必传
 * @return grpc.UnaryServerInterceptor
 * @author huwl
 */
func PeriodLimiterInterceptor(svcCtx *svc.ServiceContext) grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
		// 创建redis
		rc := redis.RedisConf{}
		switch svcCtx.Config.Env {
		case "dev":
			rc.Host = svcCtx.Config.RedisConf.Dev.Single.Host
			rc.Pass = svcCtx.Config.RedisConf.Dev.Single.Password
		case "test":
			rc.Host = svcCtx.Config.RedisConf.Test.Single.Host
			rc.Pass = svcCtx.Config.RedisConf.Test.Single.Password
		case "prod":
			rc.Host = svcCtx.Config.RedisConf.Prod.Single.Host
			rc.Pass = svcCtx.Config.RedisConf.Prod.Single.Password
		}

		if svcCtx.Config.RedisConf.Mode == "cluster" {
			rc.Type = "cluster"
		} else {
			rc.Type = "node"
		}

		// 创建全局限流器
		limiter := limit.NewPeriodLimit(
			svcCtx.Config.LimiterConf.GlobalLimiter.Period,
			svcCtx.Config.LimiterConf.GlobalLimiter.Quota,
			redis.MustNewRedis(rc),
			"global_limit")

		key := info.FullMethod
		code, err := limiter.Take(key)
		if err != nil {
			return nil, status.Error(codes.ResourceExhausted, err.Error())
		}

		// 对特定方法进行限流
		if utils.IsContainStr(svcCtx.Config.LimitFuncConf, info.FullMethod) {
			if code != limit.Allowed {
				return nil, status.Error(codes.ResourceExhausted, "服务繁忙，请稍后重试")
			}
		}

		if code != limit.Allowed {
			// 被限流，返回错误
			return nil, status.Error(codes.ResourceExhausted, "服务繁忙，请稍后重试")
		}

		return handler(ctx, req)
	}
}

/**
 * token限流器
 * 	1.以恒定速率向桶中添加令牌，请求获取到令牌才可通过
 * 	2.限制平均速率，同时允许一定程度的突发流量
 * 	3.允许流量突发，更具弹性
 * 	4.如果希望系统在空闲时能"攒下"一些容量，以应对短暂的流量峰值，比如网关、API入口等场景，令牌桶 (tokenlimit) 是更好的选择
 *
 * @param *svc.ServiceContext svcCtx -------------------------必传
 * @return grpc.UnaryServerInterceptor
 * @author huwl
 */
func TokenLimiterInterceptor(svcCtx *svc.ServiceContext) grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
		// 组装redis配置
		rc := redis.RedisConf{}
		switch svcCtx.Config.Env {
		case "dev":
			rc.Host = svcCtx.Config.RedisConf.Dev.Single.Host
			rc.Pass = svcCtx.Config.RedisConf.Dev.Single.Password
		case "test":
			rc.Host = svcCtx.Config.RedisConf.Test.Single.Host
			rc.Pass = svcCtx.Config.RedisConf.Test.Single.Password
		case "prod":
			rc.Host = svcCtx.Config.RedisConf.Prod.Single.Host
			rc.Pass = svcCtx.Config.RedisConf.Prod.Single.Password
		}

		if svcCtx.Config.RedisConf.Mode == "cluster" {
			rc.Type = "cluster"
		} else {
			rc.Type = "node"
		}

		// 创建令牌桶限流器
		// 参数：每秒生成的令牌数(rate)，令牌桶最大容量(burst)
		limiter := limit.NewTokenLimiter(
			svcCtx.Config.LimiterConf.GlobalLimiter.Rate,
			svcCtx.Config.LimiterConf.GlobalLimiter.Burst,
			redis.MustNewRedis(rc),
			"global_limit")

		// 对特定方法进行限流
		if utils.IsContainStr(svcCtx.Config.LimitFuncConf, info.FullMethod) {
			if !limiter.Allow() {
				return nil, status.Error(codes.ResourceExhausted, "服务繁忙，请稍后重试")
			}
		}
		if !limiter.Allow() {
			// 被限流，返回错误
			return nil, status.Error(codes.ResourceExhausted, "服务繁忙，请稍后重试")
		}

		return handler(ctx, req)
	}
}
